import torch
from scipy.linalg import sqrtm
import numpy as np
import time


def loss_ddpm(model,data):
    # 数据解包
    i, lv, gv, bv, av1, av2  = data
    x=torch.cat((lv, gv, av1), 1) #被生成的向量
    loss_value = model(x, bv, av2) #引入条件，计算损失
    return loss_value

def Frechet_inception_distance(real_vectors, generated_vectors):
    """
    计算两组向量之间的FID。
    
    参数:
        real_vectors (np.ndarray): 真实向量，形状为 (N, D)。
        generated_vectors (np.ndarray): 生成向量，形状为 (N, D)。
    
    返回:
        fid (float): FID值。
    """
    # 计算均值和协方差
    mu_real, sigma_real = np.mean(real_vectors, axis=0), np.cov(real_vectors, rowvar=False)
    mu_gen, sigma_gen = np.mean(generated_vectors, axis=0), np.cov(generated_vectors, rowvar=False)

    # 计算均值差异的平方
    diff = mu_real - mu_gen
    diff_squared = np.dot(diff, diff)

    # 计算协方差矩阵的平方根
    sqrt_sigma_real_gen = sqrtm(np.dot(sigma_real, sigma_gen))

    # 确保平方根是实数（避免数值误差）
    if np.iscomplexobj(sqrt_sigma_real_gen):
        sqrt_sigma_real_gen = sqrt_sigma_real_gen.real

    # 计算FID
    fid = diff_squared + np.trace(sigma_real + sigma_gen - 2 * sqrt_sigma_real_gen)
    return fid

def inception_score(features):
    """
    计算 Inception Score (IS)，用于评估生成图像的质量和多样性。

    参数:
    features : numpy.ndarray
        输入的特征向量，形状为 (N, D)，其中 N 是样本数，D 是特征维度。
        通常通过预训练的 Inception 模型提取。

    返回:
    IS : float
        Inception Score，数值越高表示生成图像的质量和多样性越好。
    """
    # 数值稳定的 softmax
    def softmax(x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

    # 避免 log(0)
    def safe_log(x):
        return np.log(np.maximum(x, 1e-10))

    # Step 1: 计算条件概率分布 p(y|x)
    p_y_given_x = softmax(features)  # 形状 (N, D)

    # Step 2: 计算边缘概率分布 p(y)
    p_y = np.mean(p_y_given_x, axis=0)  # 形状 (D,)

    # Step 3: 计算 KL 散度
    kl_divergence = p_y_given_x * (safe_log(p_y_given_x) - safe_log(p_y))  # 形状 (N, D)
    kl_divergence = np.sum(kl_divergence, axis=1)  # 形状 (N,)

    # Step 4: 计算 IS
    IS = np.exp(np.mean(kl_divergence))
    return IS